#pragma once
#include <atomic>
#include <stdint.h>
#include <stdlib.h>
#include <errno.h>
#include <string.h>
#include <condition_variable>
#include "log/log.h"

class MemPoolQueue {
public:
    MemPoolQueue(uint32_t block_size, uint32_t block_amou) :
            m_block_size(block_size), m_block_amou(block_amou) {
        m_free_block_head = nullptr;
        m_work_block_head = nullptr;
        m_work_block_tail = nullptr;

        for (int i = 0; i < m_block_amou; ++i) {
            mem_block_t* block = (mem_block_t*)malloc(sizeof(mem_block_t));
            uint8_t* buff = (uint8_t*)malloc(m_block_size);
            if (buff != nullptr && block != nullptr) {  // 加入空闲队列
                block->buff = buff;
                block->next = m_free_block_head;
                m_free_block_head = block;
            } else {
                log_error("malloc failed, %s.", strerror(errno));
                exit(EXIT_FAILURE);
            }
        }
        log_info_q(QUIET_LOG_INFO, "MemPool init seccess, block_size: %d, block_num: %d.", m_block_size, m_block_amou);
    }

    ~MemPoolQueue() {
        int free_block_size = 0;

        stop();

        while (m_free_block_head != nullptr) {
            mem_block_t* tmp_block = m_free_block_head;
            m_free_block_head = m_free_block_head->next;
            free(tmp_block->buff); tmp_block->buff = nullptr;
            free(tmp_block); tmp_block = nullptr;
            ++free_block_size;
        }

        while (m_work_block_head != nullptr) {
            mem_block_t* tmp_block = m_work_block_head;
            m_work_block_head = m_work_block_head->next;
            free(tmp_block->buff); tmp_block->buff = nullptr;
            free(tmp_block); tmp_block = nullptr;
            ++free_block_size;
        }
        if (m_block_amou == free_block_size) {
            log_info_q(QUIET_LOG_INFO, "MemPool deinit seccess, free_block_size: %d.", free_block_size);
        } else {
            log_error("MemPool deinit error, malloc block amount %d, but free block amount %d.", m_block_amou, free_block_size);
        }
    }

    int push(uint64_t pts, uint8_t* data, uint32_t len, bool need_wait) {
        std::unique_lock<std::mutex> locker(m_mutex);
        if (len > m_block_size) {
            log_error("len error, m_block_size: %d, len: %d.", m_block_size, len);
            return -1;
        }
        
        if (m_free_block_head == nullptr && need_wait == false) {
            log_error("put failed, free block is null.");
            return -1;
        }
        m_not_empty.wait(locker, [this]{return (m_need_stop || m_free_block_head != nullptr);});

        if (m_need_stop) {
            return 1;  // Stop
        }

        memcpy(m_free_block_head->buff, data, len);
        m_free_block_head->len = len;
        m_free_block_head->pts = pts;

        // block转移到工作区:插入到链表尾部
        if (m_work_block_head == nullptr) {
            mem_block_t* tmp_block = m_free_block_head;
            m_free_block_head = m_free_block_head->next;
            tmp_block->next = nullptr;
            m_work_block_head = tmp_block;
            m_work_block_tail = tmp_block;
        } else {
            mem_block_t* tmp_block = m_free_block_head;
            m_free_block_head = m_free_block_head->next;
            tmp_block->next = nullptr;
            m_work_block_tail->next = tmp_block;
            m_work_block_tail = m_work_block_tail->next;
        }
        m_not_full.notify_one();
        return 0;
    }

    uint32_t pull(uint64_t* pts, uint8_t* data, uint32_t len, bool need_wait) {
        std::unique_lock<std::mutex> locker(m_mutex);

        if (m_work_block_head == nullptr && need_wait == false) {
            log_error("take failed, work block is null.");
            return -1;
        }
        
        m_not_full.wait(locker, [this]{return (m_need_stop || m_work_block_head != nullptr);});
        if (m_need_stop) {
            return -2;  // Stop
        }

        if (len < m_work_block_head->len) {
            log_error("data len too small.");
            return -1;
        }
        memcpy(data, m_work_block_head->buff, m_work_block_head->len);
        *pts = m_work_block_head->pts;
        
        mem_block_t* tmp_block = m_work_block_head;
        m_work_block_head = m_work_block_head->next;
        
        tmp_block->next = m_free_block_head;
        m_free_block_head = tmp_block;
        m_not_empty.notify_one();
        return tmp_block->len;
    }

    uint32_t pull(bool need_wait, std::function<void(uint64_t pts, uint8_t* data, uint32_t len)> callback) {
        std::unique_lock<std::mutex> locker(m_mutex);

        if (m_work_block_head == nullptr && need_wait == false) {
            log_error("take failed, work block is null.");
            return -1;
        }
        
        m_not_full.wait(locker, [this]{return (m_need_stop || m_work_block_head != nullptr);});
        if (m_need_stop) {
            return -2;  // Stop
        }

        callback(m_work_block_head->pts, m_work_block_head->buff, m_work_block_head->len);
        
        mem_block_t* tmp_block = m_work_block_head;
        m_work_block_head = m_work_block_head->next;
        
        tmp_block->next = m_free_block_head;
        m_free_block_head = tmp_block;
        m_not_empty.notify_one();
        return tmp_block->len;
    }

    int size(uint32_t* free_block_amount, uint32_t* work_block_amount) {
        std::unique_lock<std::mutex> locker(m_mutex);
        mem_block_t* tmp_block;

        if (free_block_amount != nullptr) {
            tmp_block = m_free_block_head;
            *free_block_amount = 0;
            while (tmp_block != nullptr) {
                ++(*free_block_amount);
                tmp_block = tmp_block->next;
            }
        }

        if (work_block_amount != nullptr) {
            tmp_block = m_work_block_head;
            *work_block_amount = 0;
            while (tmp_block != nullptr) {
                ++(*work_block_amount);
                tmp_block = tmp_block->next;
            }
        }
        
        return 0;
    }

    int stop() {
        {
            std::lock_guard<std::mutex> locker(m_mutex);
            m_need_stop = true;
        }
        m_not_empty.notify_all();
        m_not_full.notify_all();
        return 0;
    }

private:
    // int get_work_head_block() {
    //     std::unique_lock<std::mutex> locker(m_mutex);

    //     return 0;
    // }

    // int get_work_tail_block() {
    //     return 0;
    // }

    // int put_work_head_block() {
    //     return 0;
    // }

    // int put_work_tail_block() {
    //     return 0;
    // }

    // int get_free_head_block() {
    //     return 0;
    // }

    // int get_free_tail_block() {
    //     return 0;
    // }

    // int put_free_head_block() {
    //     return 0;
    // }

    // int put_free_tail_block() {
    //     return 0;
    // }

private:
    /* logger config */
    const int QUIET_LOG_DEBUG{1};
    const int QUIET_LOG_INFO{1};

    typedef struct mem_block_t {
        uint32_t len;       // 有效数据长度
        uint8_t* buff;      // 挂载的内存地址
        mem_block_t* next;  // 下一个块的地址
        // 拓展数据
        uint64_t pts;       // 时间戳
    } mem_block_t;

    uint32_t m_block_size;
    uint32_t m_block_amou;
    
    std::atomic_bool m_need_stop{false};
    std::mutex m_mutex;  // 互斥量和条件变量结合起来使用
    std::condition_variable m_not_empty;  // 空闲块不为空的条件变量
    std::condition_variable m_not_full;  // 空闲块没有满的条件变量

    mem_block_t* m_free_block_head{};
    mem_block_t* m_work_block_head{};
    mem_block_t* m_work_block_tail{};
};